"""
@description: Handles loading of pre-trained models and tokenizers, including quantization logic.
@author: yangkw
@date: 2025-05-02
"""
# -*- coding: utf-8 -*-
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig
)
import logging
from .utils import ModelArguments, QuantizationArguments # 从同级目录的 utils 导入

logger = logging.getLogger(__name__)

def load_model_and_tokenizer(model_args: ModelArguments, quant_args: QuantizationArguments):
    """
    加载预训练模型和分词器。

    根据提供的参数处理 4-bit 量化加载。

    Args:
        model_args (ModelArguments): 包含模型名称/路径等信息的参数对象。
        quant_args (QuantizationArguments): 包含量化配置的参数对象。

    Returns:
        tuple: 包含加载的模型 (model) 和分词器 (tokenizer) 的元组。
    """
    logger.info(f"开始加载模型: {model_args.model_name_or_path}")

    # --- 1. 加载 Tokenizer ---
    tokenizer_name = model_args.tokenizer_name_or_path if model_args.tokenizer_name_or_path else model_args.model_name_or_path
    try:
        # 考虑添加 use_fast=True? 对于某些模型可能需要设置为 False
        # trust_remote_code=True 对于某些社区模型可能是必需的
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
        logger.info(f"成功加载 Tokenizer: {tokenizer_name}")

        # 设置 padding token (如果模型没有定义)
        # 对于 Llama 等模型，通常建议将 pad_token 设置为 eos_token
        if tokenizer.pad_token is None:
            if tokenizer.eos_token is not None:
                tokenizer.pad_token = tokenizer.eos_token
                logger.info(f"Tokenizer 没有 pad_token，已设置为 eos_token: {tokenizer.eos_token}")
            else:
                # 如果连 eos_token 都没有，添加一个新的 pad token 可能需要调整模型嵌入层大小
                # 这里暂时只记录警告
                logger.warning("Tokenizer 既没有 pad_token 也没有 eos_token。可能需要在后续处理中指定或添加。")
                # tokenizer.add_special_tokens({'pad_token': '[PAD]'})
                # model.resize_token_embeddings(len(tokenizer)) # 如果添加了新 token

    except Exception as e:
        logger.error(f"加载 Tokenizer '{tokenizer_name}' 失败: {e}")
        raise

    # --- 2. 配置量化 (如果启用) ---
    quantization_config = None
    compute_dtype = None # 用于 BitsAndBytesConfig 的计算类型

    if quant_args.load_in_4bit:
        logger.info("启用 4-bit 量化加载。")
    else:
        logger.info("未启用 4-bit 量化加载。")
        # 如果不量化，可以根据训练参数决定使用 bf16 还是 fp16 (但这通常在 Trainer 或 Accelerate 中处理)
        # compute_dtype = torch.bfloat16 if training_args.bf16 else (torch.float16 if training_args.fp16 else torch.float32)

  # 确定计算数据类型
    if quant_args.bnb_4bit_compute_dtype == "bfloat16":
        # 检查当前环境是否支持 bfloat16
        if torch.cuda.is_bf16_supported():
            compute_dtype = torch.bfloat16
            logger.info("使用 bfloat16 作为 4-bit 计算数据类型。")
        else:
            logger.warning("当前环境不支持 bfloat16，将回退到 float16。")
            compute_dtype = torch.float16
    elif quant_args.bnb_4bit_compute_dtype == "float16":
        compute_dtype = torch.float16
        logger.info("使用 float16 作为 4-bit 计算数据类型。")
    else:
        logger.warning(f"未知的 bnb_4bit_compute_dtype: {quant_args.bnb_4bit_compute_dtype}。将使用默认值。")
        # BitsAndBytesConfig 内部可能有默认处理，或者这里可以设置一个默认值
        compute_dtype = torch.float16 # 假设默认

    try:
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=quant_args.load_in_4bit,
            bnb_4bit_quant_type=quant_args.bnb_4bit_quant_type, # "nf4" 或 "fp4"
            bnb_4bit_compute_dtype=compute_dtype,
            bnb_4bit_use_double_quant=quant_args.bnb_4bit_use_double_quant,
        )
        logger.info(f"BitsAndBytesConfig 配置完成: {quantization_config}")
    except ImportError:
        logger.error("bitsandbytes 库未安装或版本不兼容，无法进行 4-bit 量化。请安装 'pip install bitsandbytes'")
        raise
    except Exception as e:
        logger.error(f"创建 BitsAndBytesConfig 时出错: {e}")
        raise

    # --- 3. 加载模型 ---
    try:
        # device_map="auto" 会尝试自动将模型分片到可用 GPU 和 CPU，减少单卡显存压力
        # Accelerate 会进一步优化设备映射，尤其是在多 GPU 场景下
        # 对于 DeepSpeed ZeRO Stage 3，device_map 可能需要设置为 None 或特定值，由 DeepSpeed 控制
        # 在使用 Accelerate + DeepSpeed 时，通常建议不设置 device_map，让 Accelerate 处理
        # 关闭flash attention以节省显存
        use_flash_attention = False
        model = AutoModelForCausalLM.from_pretrained(
            model_args.model_name_or_path,
            quantization_config=quantization_config, # 应用量化配置
            torch_dtype=compute_dtype if quantization_config else None, # 如果量化，dtype 由 BNB 控制；否则可以指定
            trust_remote_code=True, # 对于某些模型是必需的
            # device_map="auto" # 通常由 Accelerate 处理，这里注释掉
            # - sdpa: 显存不足时的默认选项
            attn_implementation="flash_attention_2" if use_flash_attention else "sdpa",
            use_cache=False,           # 关闭KV缓存以节省显存
        )
        logger.info(f"成功加载模型: {model_args.model_name_or_path}")

        # 如果使用了量化，打印模型占用的内存（近似值）
        if quant_args.load_in_4bit:
             try:
                 model_memory_footprint = model.get_memory_footprint()
                 logger.info(f"量化模型加载后的内存占用 (近似): {model_memory_footprint / (1024**3):.2f} GB")
             except AttributeError:
                 logger.warning("无法获取量化模型的内存占用信息 (get_memory_footprint 方法可能不存在)。")


    except Exception as e:
        logger.error(f"加载模型 '{model_args.model_name_or_path}' 失败: {e}")
        raise

    return model, tokenizer

# 示例用法 (通常在主脚本中调用)
if __name__ == '__main__':
    from .utils import parse_args_from_yaml
    from pathlib import Path

    # 假设 config 文件在父目录的 config/ 下
    script_dir = Path(__file__).parent
    project_root = script_dir.parent
    example_config_path = project_root / "config" / "training_args.yaml"

    if example_config_path.exists():
        try:
            model_args, _, _, quant_args, _, _ = parse_args_from_yaml(example_config_path)

            # 为了测试，强制启用量化 (如果配置文件没开)
            # quant_args.load_in_4bit = True

            logger.info("开始测试模型和 Tokenizer 加载...")
            model, tokenizer = load_model_and_tokenizer(model_args, quant_args)
            logger.info("模型和 Tokenizer 加载测试成功!")

            # 打印模型和 tokenizer 信息 (可选)
            # print("\nModel:", model)
            # print("\nTokenizer:", tokenizer)
            print(f"\nModel class: {model.__class__.__name__}")
            print(f"Tokenizer class: {tokenizer.__class__.__name__}")
            if hasattr(model, 'config'):
                 print(f"Model config: {model.config._name_or_path}")

        except (FileNotFoundError, ValueError, TypeError, RuntimeError, ImportError) as e:
            logger.error(f"模型加载测试失败: {e}")
        except Exception as e:
             logger.error(f"发生未知错误: {e}")
    else:
        logger.warning(f"示例配置文件未找到: {example_config_path}，无法执行模型加载测试。")
